import torch
import torch.nn as nn
from   diffusers import AutoencoderKL

class Resnet(nn.Module):
    def __init__(self, dim_in, dim_out):
        super().__init__()

        self.s = torch.nn.Sequential(
            torch.nn.GroupNorm(num_groups=32, num_channels=dim_in, eps=1e-6, affine=True),
            torch.nn.SiLU(),
            torch.nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1),
            torch.nn.GroupNorm(num_groups=32, num_channels=dim_out, eps=1e-6, affine=True),
            torch.nn.SiLU(),
            torch.nn.Conv2d(dim_out, dim_out, kernel_size=3, stride=1, padding=1),
        )

        self.res = None
        if dim_in != dim_out:
            self.res = torch.nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=1, padding=0)

    def forward(self, x):             # x -> [1, 128, 10, 10]
        res = x
        if self.res:
            res = self.res(x)         # [1, 128, 10, 10] -> [1, 256, 10, 10]
        return res + self.s(x)        # [1, 128, 10, 10] -> [1, 256, 10, 10]
    

class Attention(nn.Module):
    def __init__(self):
        super(Attention, self).__init__()
        self.norm = nn.GroupNorm(num_channels=512, num_groups=32,
                                eps=1e-6, affine=True)
        self.Q = torch.nn.Linear(512, 512)
        self.K = torch.nn.Linear(512, 512)
        self.V = torch.nn.Linear(512, 512)
        self.O = torch.nn.Linear(512, 512)
        
    
    def forward(self, x):                  # x: [1, 512, 64, 64]
        res = x
        x   = self.norm(x)                 # x: [1, 512, 64, 64]
        
        # [1, 512, 64, 64] -> [1, 512, 4096] -> [1, 4096, 512]
        x   = x.flatten(start_dim=2).transpose(1, 2)
        
        q   = self.Q(x)                    # 线性运算,维度不变 [1, 4096, 512]
        k   = self.K(x)
        v   = self.V(x)
        
        k   = k.transpose(1, 2)            # [1, 4096, 512] -> [1, 512, 4096]
        
        # [1, 4096, 512] * [1, 512, 4096] -> [1, 4096, 4096]
        #0.044194173824159216 = 1 / 512**0.5
        #atten = q.bmm(k) * 0.044194173824159216

        #照理来说应该是等价的,但是却有很小的误差
        atten = torch.baddbmm(torch.empty(1, 4096, 4096, device=q.device),
                              q, k, beta=0, alpha=0.044194173824159216)
        atten = torch.softmax(atten, dim=-1)
        atten = atten.bmm(v)               # [1, 4096, 4096]*[1, 4096, 512]->[1, 4096, 512]
        atten = self.O(atten)              # 线性运算,维度不变 [1, 4096, 512]
        
        # [1, 4096, 512] -> [1, 512, 4096] -> [1, 512, 64, 64]
        atten = atten.transpose(1, 2).reshape(-1, 512, 64, 64)
        return atten + res                 # shortcut connection
    
    
class Pad(nn.Module):
    def forward(self, x):
        """
        @brief nn.functional.pad对张量进行填充
        
        @param input (Tensor):          要填充的输入张量。
        @param pad   (tuple):           一个元组，表示每个维度两侧的填充量。元组的长度应该是输入维度的两倍。例如，对于一
                                        个二维输入（如图像），pad 应该有四个元素 (left, right, top, bottom)，这表
                                        示左侧填充left个单位，右侧填充right个单位，顶部填充top个单位，底部填充bottom
                                        个单位。
        @param mode  (str, optional):   填充的模式。默认是 'constant'，表示用常数值填充。其他选项包括 'reflect'（边
                                        界反射填充），'replicate'（边界复制填充）和 'circular'（循环填充）。每种模式
                                        的填充方式略有不同。
        @param value (float, optional): 当mode='constant'时，用于填充的常数值。默认是0。
        """
        return nn.functional.pad(x, (0,1,0,1), mode="constant", value=0)
    

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.encoder = nn.Sequential(
            torch.nn.Conv2d(3, 128, kernel_size=3, stride=1, padding=1),   # in
            torch.nn.Sequential(                                           # down
                Resnet(128, 128), Resnet(128, 128),
                torch.nn.Sequential(Pad(),
                    torch.nn.Conv2d(128, 128, 3, stride=2, padding=0),
                ),
            ),
            torch.nn.Sequential( Resnet(128, 256), Resnet(256, 256),
                torch.nn.Sequential(Pad(),
                    torch.nn.Conv2d(256, 256, 3, stride=2, padding=0),
                ),
            ),
            torch.nn.Sequential( Resnet(256, 512), Resnet(512, 512),
                torch.nn.Sequential(Pad(),
                    torch.nn.Conv2d(512, 512, 3, stride=2, padding=0),
                ),
            ),
            torch.nn.Sequential(Resnet(512, 512),Resnet(512, 512)),

            # mid
            torch.nn.Sequential( Resnet(512, 512), Attention(), Resnet(512, 512)),

            # out
            torch.nn.Sequential(
                torch.nn.GroupNorm(num_channels=512, num_groups=32, eps=1e-6),
                torch.nn.SiLU(), torch.nn.Conv2d(512, 8, 3, padding=1),
            ),
            torch.nn.Conv2d(8, 8, 1),    # 正态分布层
        )
        
        self.decoder = nn.Sequential(
            torch.nn.Conv2d(4, 4, 1),    # 正态分布层
            torch.nn.Conv2d(4, 512, kernel_size=3, stride=1, padding=1),  # in

            # middle
            torch.nn.Sequential(Resnet(512, 512), Attention(), Resnet(512, 512)),

            # up
            torch.nn.Sequential(Resnet(512, 512), Resnet(512, 512), Resnet(512, 512),
                torch.nn.Upsample(scale_factor=2.0, mode='nearest'),
                torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),
            ),
            torch.nn.Sequential(Resnet(512, 512), Resnet(512, 512), Resnet(512, 512),
                torch.nn.Upsample(scale_factor=2.0, mode='nearest'),
                torch.nn.Conv2d(512, 512, kernel_size=3, padding=1),
            ),
            torch.nn.Sequential(Resnet(512, 256), Resnet(256, 256), Resnet(256, 256),
                torch.nn.Upsample(scale_factor=2.0, mode='nearest'),
                torch.nn.Conv2d(256, 256, kernel_size=3, padding=1),
            ),
            torch.nn.Sequential(Resnet(256, 128), Resnet(128, 128), Resnet(128, 128)),

            # out
            torch.nn.Sequential(
                torch.nn.GroupNorm(num_channels=128, num_groups=32, eps=1e-6),
                torch.nn.SiLU(), torch.nn.Conv2d(128, 3, 3, padding=1),
            )
        )
        
    def sample(self, h):                 # h -> [1, 8, 64, 64]
        mean   = h[:,:4]
        logvar = h[:,4:]
        std    = logvar.exp()**0.5
        h      = torch.randn(mean.shape, device=mean.device)
        h      = mean + std * h
        return h
    
    def forward(self, x):                # x -> [1, 3, 512, 512]
        h = self.encoder(x)              # 编码: [1, 3, 512, 512] -> [1, 8, 64, 64]
        h = self.sample(h)               # 投影: [1, 8, 64, 64]   -> [1, 4, 64, 64]
        h = self.decoder(h)              # 解码: [1, 4, 64, 64]   -> [1, 3, 512, 512]
        return h
           
def build_vae():
    def load_res(model, param):
        model.s[0].load_state_dict(param.norm1.state_dict())
        model.s[2].load_state_dict(param.conv1.state_dict())
        model.s[3].load_state_dict(param.norm2.state_dict())
        model.s[5].load_state_dict(param.conv2.state_dict())

        if isinstance(model.res, torch.nn.Module):
            model.res.load_state_dict(param.conv_shortcut.state_dict())

    def load_atten(model, param):
        model.norm.load_state_dict(param.group_norm.state_dict())
        model.q.load_state_dict(param.query.state_dict())
        model.k.load_state_dict(param.key.state_dict())
        model.v.load_state_dict(param.value.state_dict())
        model.out.load_state_dict(param.proj_attn.state_dict())
    
    # 加载预训练模型的参数
    params = AutoencoderKL.from_pretrained('lansinuote/diffsion_from_scratch.params', subfolder='vae')

    vae = VAE()
    vae.encoder[0].load_state_dict(params.encoder.conv_in.state_dict())   # encoder.in
    
    for i in range(4):                                                    # encoder.down
        load_res(vae.encoder[i + 1][0], params.encoder.down_blocks[i].resnets[0])
        load_res(vae.encoder[i + 1][1], params.encoder.down_blocks[i].resnets[1])

        if i != 3:
            vae.encoder[i + 1][2][1].load_state_dict(
                params.encoder.down_blocks[i].downsamplers[0].conv.state_dict())
            
    # encoder.mid
    load_res(vae.encoder[5][0], params.encoder.mid_block.resnets[0])
    load_res(vae.encoder[5][2], params.encoder.mid_block.resnets[1])
    load_atten(vae.encoder[5][1], params.encoder.mid_block.attentions[0])

    # encoder.out
    vae.encoder[6][0].load_state_dict(params.encoder.conv_norm_out.state_dict())
    vae.encoder[6][2].load_state_dict(params.encoder.conv_out.state_dict())

    # encoder.正态分布层
    vae.encoder[7].load_state_dict(params.quant_conv.state_dict())

    #decoder.正态分布层
    vae.decoder[0].load_state_dict(params.post_quant_conv.state_dict())

    # decoder.in
    vae.decoder[1].load_state_dict(params.decoder.conv_in.state_dict())

    # decoder.mid
    load_res(vae.decoder[2][0], params.decoder.mid_block.resnets[0])
    load_res(vae.decoder[2][2], params.decoder.mid_block.resnets[1])
    load_atten(vae.decoder[2][1], params.decoder.mid_block.attentions[0])

    # decoder.up
    for i in range(4):
        load_res(vae.decoder[i + 3][0], params.decoder.up_blocks[i].resnets[0])
        load_res(vae.decoder[i + 3][1], params.decoder.up_blocks[i].resnets[1])
        load_res(vae.decoder[i + 3][2], params.decoder.up_blocks[i].resnets[2])

        if i != 3:
            vae.decoder[i + 3][4].load_state_dict(
                params.decoder.up_blocks[i].upsamplers[0].conv.state_dict())

    # decoder.out
    vae.decoder[7][0].load_state_dict(params.decoder.conv_norm_out.state_dict())
    vae.decoder[7][2].load_state_dict(params.decoder.conv_out.state_dict())
    
    return vae